import argparse
import os
import yaml
from EntityEmbedding.script.entity_emb_main import entity_emb_main
from GraphEmbedding.script.main import graph_emb_main
from Cluster.cluster import cluster_main
from config import setup_seed, SEED, project_path
from Preprocessed.log_extract import log_extract_main
from Preprocessed.demo_pre import preprocess_train_data, preprocess_test_data
def load_yaml(yaml_path):
    with open(yaml_path, 'r', encoding='utf-8') as f:
            config = yaml.safe_load(f)
    return config

setup_seed(SEED)
# parser = argparse.ArgumentParser()
# parser.add_argument('--config',help='config.yml path')
# parser.add_argument('--dataset',help='gaia or others')
# parser.add_argument('--op',help='train or test')
# parser.add_argument('--start_time', type=int, help='')
# parser.add_argument('--end_time', type=int, help='')
# args = parser.parse_args()

def main(op: str, start_time: int, end_time: int):
    # load yaml
    DATASET = "demo"
    config_a = load_yaml(project_path/"config.yml")
    config = config_a[DATASET]
    # 提取日志模板
    log_extract_main(project_path/f"origin_data/{DATASET}")
    train_case_name = config['entity_embedding']['failure_type']
    test_case_name = ['test']
    # 预处理三种模态数据
    if op == 'train':
        preprocess_train_data(project_path/f"origin_data/{DATASET}", project_path / f'data/{DATASET}_TKG/')
        # entity embedding
        normal_data_length = entity_emb_main(config, train_case_name, test_case_name, if_train=True)
        # graph embedding:
        graph_emb_main(config, if_train=True)
    elif op == 'test':
        preprocess_test_data(project_path/f"origin_data/{DATASET}", project_path / f'data/{DATASET}_TKG/', start_time, end_time)
        # entity embedding
        case_name = ['test']
        normal_data_length = entity_emb_main(config, train_case_name, test_case_name, if_train=False)

        # graph embedding:
        graph_emb_main(config, if_train=False)

        # cluster
        cluster_main(config, normal_data_length)
    else:
        print(f"需要传入op参数，train或者test")

if __name__ == "__main__":
     main("train", 0,0)
    #  main("test", 1686268500, 1686269640)

